#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ChromaDB 自动备份工具
定期自动备份聊天记录，防止数据丢失
"""

import os
import sys
import json
import shutil
from pathlib import Path
from datetime import datetime
import schedule
import time

# 禁用遥测
os.environ['ANONYMIZED_TELEMETRY'] = 'False'

try:
    import chromadb
    from chromadb.config import Settings
except ImportError:
    print("❌ 未安装 chromadb，请运行: pip install chromadb schedule")
    sys.exit(1)


class ChromaDBBackup:
    """ChromaDB 自动备份管理器"""
    
    def __init__(self):
        self.project_root = Path(__file__).parent.parent
        self.chroma_path = self.project_root / 'chroma'
        self.backup_dir = self.project_root / 'chroma_backups'
        self.backup_dir.mkdir(exist_ok=True)
        
        # 创建备份子目录
        self.daily_backup = self.backup_dir / 'daily'
        self.weekly_backup = self.backup_dir / 'weekly'
        self.manual_backup = self.backup_dir / 'manual'
        
        for dir_path in [self.daily_backup, self.weekly_backup, self.manual_backup]:
            dir_path.mkdir(exist_ok=True)
        
        print(f"📁 备份目录: {self.backup_dir}")
    
    def connect_chromadb(self):
        """连接到 ChromaDB"""
        try:
            client = chromadb.PersistentClient(
                path=str(self.chroma_path),
                settings=Settings(
                    anonymized_telemetry=False,
                    allow_reset=True
                )
            )
            collection = client.get_collection(name="chat_history")
            return client, collection
        except Exception as e:
            print(f"❌ 连接 ChromaDB 失败: {e}")
            return None, None
    
    def export_to_json(self, collection, output_file):
        """导出数据为 JSON 格式"""
        try:
            # 获取所有记录
            results = collection.get(include=['metadatas', 'documents'])
            
            if not results['ids']:
                print("⚠️  数据库为空，无需备份")
                return False
            
            # 构建导出数据
            export_data = {
                'backup_time': datetime.now().isoformat(),
                'total_records': len(results['ids']),
                'conversations': []
            }
            
            for i in range(len(results['ids'])):
                export_data['conversations'].append({
                    'id': results['ids'][i],
                    'metadata': results['metadatas'][i] if results.get('metadatas') else {},
                    'document': results['documents'][i] if results.get('documents') else ''
                })
            
            # 写入 JSON 文件
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(export_data, f, ensure_ascii=False, indent=2)
            
            return True
        except Exception as e:
            print(f"❌ 导出 JSON 失败: {e}")
            return False
    
    def backup_database_files(self, backup_type='daily'):
        """备份整个数据库文件夹"""
        if not self.chroma_path.exists():
            print("❌ 数据库文件夹不存在")
            return False
        
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        
        # 根据备份类型选择目录
        if backup_type == 'daily':
            backup_root = self.daily_backup
        elif backup_type == 'weekly':
            backup_root = self.weekly_backup
        else:
            backup_root = self.manual_backup
        
        backup_folder = backup_root / f"chroma_{timestamp}"
        
        try:
            # 复制整个数据库文件夹
            shutil.copytree(self.chroma_path, backup_folder)
            print(f"✅ 数据库文件备份成功: {backup_folder.name}")
            return True
        except Exception as e:
            print(f"❌ 备份失败: {e}")
            return False
    
    def create_full_backup(self, backup_type='daily'):
        """创建完整备份（JSON + 数据库文件）"""
        print(f"\n{'='*60}")
        print(f"🔄 开始 {backup_type} 备份...")
        print(f"{'='*60}")
        
        # 连接数据库
        client, collection = self.connect_chromadb()
        if not collection:
            return False
        
        count = collection.count()
        print(f"📊 当前有 {count} 条记录")
        
        if count == 0:
            print("⚠️  数据库为空，跳过备份")
            return False
        
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        
        # 根据备份类型选择目录
        if backup_type == 'daily':
            backup_root = self.daily_backup
        elif backup_type == 'weekly':
            backup_root = self.weekly_backup
        else:
            backup_root = self.manual_backup
        
        # 1. 导出 JSON
        json_file = backup_root / f"backup_{timestamp}.json"
        if self.export_to_json(collection, json_file):
            file_size = json_file.stat().st_size / 1024  # KB
            print(f"✅ JSON 备份成功: {json_file.name} ({file_size:.1f} KB)")
        
        # 2. 备份数据库文件
        self.backup_database_files(backup_type)
        
        print(f"✅ {backup_type.upper()} 备份完成！")
        print(f"{'='*60}\n")
        
        return True
    
    def cleanup_old_backups(self, backup_type='daily', keep_count=7):
        """清理旧备份"""
        if backup_type == 'daily':
            backup_root = self.daily_backup
        elif backup_type == 'weekly':
            backup_root = self.weekly_backup
        else:
            return
        
        # 获取所有备份文件和文件夹
        json_files = sorted(backup_root.glob('backup_*.json'), reverse=True)
        db_folders = sorted(backup_root.glob('chroma_*'), reverse=True)
        
        # 只保留最新的 N 个备份
        for old_file in json_files[keep_count:]:
            old_file.unlink()
            print(f"🗑️  删除旧备份: {old_file.name}")
        
        for old_folder in db_folders[keep_count:]:
            shutil.rmtree(old_folder)
            print(f"🗑️  删除旧备份: {old_folder.name}")
    
    def list_backups(self):
        """列出所有备份"""
        print("\n" + "="*60)
        print("📋 备份列表")
        print("="*60)
        
        for backup_type, backup_dir in [
            ('每日备份', self.daily_backup),
            ('每周备份', self.weekly_backup),
            ('手动备份', self.manual_backup)
        ]:
            print(f"\n{backup_type} ({backup_dir.name}):")
            json_files = sorted(backup_dir.glob('backup_*.json'), reverse=True)
            
            if not json_files:
                print("  （暂无备份）")
                continue
            
            for i, backup_file in enumerate(json_files[:5], 1):  # 只显示最新5个
                try:
                    with open(backup_file, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                    backup_time = data.get('backup_time', 'Unknown')
                    record_count = data.get('total_records', 0)
                    file_size = backup_file.stat().st_size / 1024
                    
                    print(f"  {i}. {backup_file.name}")
                    print(f"     时间: {backup_time[:19]}")
                    print(f"     记录数: {record_count} 条")
                    print(f"     大小: {file_size:.1f} KB")
                except:
                    print(f"  {i}. {backup_file.name} (无法读取)")
        
        print("\n" + "="*60 + "\n")


def manual_backup():
    """手动备份"""
    backup_manager = ChromaDBBackup()
    backup_manager.create_full_backup('manual')


def daily_backup():
    """每日自动备份"""
    backup_manager = ChromaDBBackup()
    backup_manager.create_full_backup('daily')
    backup_manager.cleanup_old_backups('daily', keep_count=7)  # 保留7天


def weekly_backup():
    """每周备份"""
    backup_manager = ChromaDBBackup()
    backup_manager.create_full_backup('weekly')
    backup_manager.cleanup_old_backups('weekly', keep_count=4)  # 保留4周


def start_auto_backup_service():
    """启动自动备份服务"""
    print("="*60)
    print("🤖 ChromaDB 自动备份服务")
    print("="*60)
    print()
    print("📅 备份计划:")
    print("  - 每日备份: 每天凌晨 2:00")
    print("  - 每周备份: 每周日凌晨 3:00")
    print("  - 保留策略: 每日备份保留7天，每周备份保留4周")
    print()
    print("💡 提示: 按 Ctrl+C 停止服务")
    print("="*60)
    print()
    
    # 启动时立即备份一次
    print("🔄 启动时执行首次备份...")
    manual_backup()
    
    # 设置定时任务
    schedule.every().day.at("02:00").do(daily_backup)
    schedule.every().sunday.at("03:00").do(weekly_backup)
    
    print("✅ 自动备份服务已启动！")
    print()
    
    try:
        while True:
            schedule.run_pending()
            time.sleep(60)  # 每分钟检查一次
    except KeyboardInterrupt:
        print("\n\n👋 自动备份服务已停止")


def restore_from_backup(backup_file):
    """从备份恢复数据"""
    backup_manager = ChromaDBBackup()
    
    print("="*60)
    print("🔄 从备份恢复数据")
    print("="*60)
    print()
    
    # 读取备份文件
    if not os.path.exists(backup_file):
        print(f"❌ 备份文件不存在: {backup_file}")
        return False
    
    print(f"📂 读取备份文件: {backup_file}")
    
    try:
        with open(backup_file, 'r', encoding='utf-8') as f:
            backup_data = json.load(f)
    except Exception as e:
        print(f"❌ 读取备份文件失败: {e}")
        return False
    
    conversations = backup_data.get('conversations', [])
    print(f"✅ 找到 {len(conversations)} 条记录")
    
    # 连接数据库
    client, collection = backup_manager.connect_chromadb()
    if not collection:
        return False
    
    # 恢复数据
    success_count = 0
    skip_count = 0
    error_count = 0
    
    for i, conv in enumerate(conversations):
        try:
            # 检查是否已存在
            try:
                existing = collection.get(ids=[conv['id']])
                if existing and existing['ids']:
                    skip_count += 1
                    continue
            except:
                pass
            
            # 添加记录
            collection.add(
                ids=[conv['id']],
                documents=[conv['document']],
                metadatas=[conv['metadata']]
            )
            
            success_count += 1
            print(f"  ✅ [{i+1}/{len(conversations)}] 恢复成功")
            
        except Exception as e:
            error_count += 1
            print(f"  ❌ [{i+1}/{len(conversations)}] 恢复失败: {e}")
    
    print()
    print("="*60)
    print(f"📊 恢复完成统计:")
    print(f"  ✅ 成功: {success_count} 条")
    print(f"  ⚠️  跳过: {skip_count} 条")
    print(f"  ❌ 失败: {error_count} 条")
    print(f"  📊 当前总记录: {collection.count()} 条")
    print("="*60)
    
    return success_count > 0


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(
        description='ChromaDB 自动备份工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法：
  python auto_backup.py --backup              # 立即创建手动备份
  python auto_backup.py --list                # 列出所有备份
  python auto_backup.py --start-service       # 启动自动备份服务
  python auto_backup.py --restore backup.json # 从备份恢复数据
        """
    )
    
    parser.add_argument('--backup', action='store_true', help='立即创建手动备份')
    parser.add_argument('--list', action='store_true', help='列出所有备份')
    parser.add_argument('--start-service', action='store_true', help='启动自动备份服务')
    parser.add_argument('--restore', type=str, help='从指定备份文件恢复数据')
    
    args = parser.parse_args()
    
    if args.backup:
        manual_backup()
    elif args.list:
        backup_manager = ChromaDBBackup()
        backup_manager.list_backups()
    elif args.start_service:
        start_auto_backup_service()
    elif args.restore:
        restore_from_backup(args.restore)
    else:
        parser.print_help()


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n👋 已取消")
    except Exception as e:
        print(f"\n❌ 错误: {e}")
        import traceback
        traceback.print_exc()
